package com.cnd.scm.common.inner.async;

import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.util.Map;

/**
 * @version V1.0
 * @ClassName: ContextCopyingDecorator
 * @Description:TODO()
 * @author: 李启联
 * @date: 2020-01-17 11:00
 * @Copyright: Copyright (c), 2020
 * 申明：本内容未经授权,禁止外泄以及用于其他的商业项目.
 */
// https://stackoverflow.com/questions/23732089/how-to-enable-request-scope-in-async-task-executor
// 传递RequestAttributes and MDC
// SecurityContext
public class ContextCopyingDecorator implements TaskDecorator {

    @Override
    public Runnable decorate(Runnable runnable) {
        try {
            RequestAttributes context = RequestContextHolder.currentRequestAttributes();
            Map<String,String> previous = MDC.getCopyOfContextMap();
            SecurityContext securityContext = SecurityContextHolder.getContext();// 1
            return () -> {
                try {
                    if(previous==null){
                        MDC.clear();
                    }else{
                        MDC.setContextMap(previous);
                    }

                    RequestContextHolder.setRequestAttributes(context);
                    SecurityContextHolder.setContext(securityContext);// 2
                    runnable.run();
                } finally {
                    RequestContextHolder.resetRequestAttributes();
                    // 清除操作
                    SecurityContextHolder.clearContext();// 3

                    if(previous==null){
                        MDC.clear();
                    }else{
                        MDC.setContextMap(previous);
                    }
                }
            };
        } catch (IllegalStateException e) {
            return runnable;
        }
    }
}